package main;

import java.util.HashMap;
import java.util.Map;

public class RandomList<T> {
    private Map<T, Integer> elements;
    private int totalWeight;

    public RandomList() {
        elements = new HashMap<>();
        totalWeight = 0;
    }

    public void add(T element, int weight) {
        elements.merge(element, weight, (oldWeight, newWeight) -> oldWeight + newWeight);
        totalWeight += weight;
    }

    public void remove(T element, int weight) {
        int toRemove = Math.min(weight, elements.get(element));
        elements.compute(element, (elem, oldWeight) -> toRemove >= oldWeight ? null : oldWeight - toRemove);
        totalWeight -= toRemove;
    }

    public T getElement() {
        long roll = RandomProvider.roll(totalWeight);
        for (Map.Entry<T, Integer> entry : elements.entrySet()) {
            if (roll <= entry.getValue()) {
                return entry.getKey();
            }

            roll -= entry.getValue();
        }

        throw new RuntimeException("Can't choose element from random list " + this);
    }
    
    public int getWeight(T element) {
        return elements.getOrDefault(element, 0);
    }
}
